package h264

import (
	"bufio"
	"bytes"
	"fmt"
	"gitee.com/frank2020/joy4/av"
	"gitee.com/frank2020/joy4/av/avutil"
	"gitee.com/frank2020/joy4/codec/h264parser"
	"gitee.com/frank2020/joy4/utils/bits/pio"
	"io"
	"log"
)

func Handler(h *avutil.RegisterHandler) {
	h.Probe = func(b []byte) bool {
		return b[0] == '0' && b[1] == '0' && b[2] == '0' && b[3] == '1'
	}

	h.Ext = ".264"

	h.ReaderDemuxer = func(r io.Reader) av.Demuxer {
		return NewDemuxer(r)
	}

	h.WriterMuxer = func(w io.Writer) av.Muxer {
		return NewMuxer(w)
	}

}

type Muxer struct {
	bufw    writeFlusher
	streams []av.CodecData
}

type writeFlusher interface {
	io.Writer
	Flush() error
}

func NewMuxerWriteFlusher(w writeFlusher) *Muxer {
	return &Muxer{
		bufw: w,
	}
}

func NewMuxer(w io.Writer) *Muxer {
	return NewMuxerWriteFlusher(bufio.NewWriterSize(w, pio.RecommendBufioSize))
}

func (m *Muxer) WriteHeader(streams []av.CodecData) (err error) {
	m.streams = streams
	for _, stream := range streams {
		if stream.Type().IsVideo() {
			vstream := stream.(av.VideoCodecData)
			codec := vstream.(h264parser.CodecData)
			m.bufw.Write([]byte{0, 0, 0, 1})
			m.bufw.Write(codec.SPS())
			m.bufw.Write([]byte{0, 0, 0, 1})
			m.bufw.Write(codec.PPS())
		}
	}
	return
}

func (m *Muxer) WritePacket(pkt av.Packet) (err error) {
	if m.streams[pkt.Idx].Type().IsVideo() {
		nalus, _ := h264parser.SplitNALUs(pkt.Data)
		for _, nalu := range nalus {
			m.bufw.Write([]byte{0, 0, 0, 1})
			m.bufw.Write(nalu)
		}

	}
	return
}

func (m *Muxer) WriteTrailer() (err error) {
	if err = m.bufw.Flush(); err != nil {
		return
	}
	return
}

type Demuxer struct {
	scanner *bufio.Scanner

	pkts chan *av.Packet

	pps []byte
	sps []byte

	gotCodecData bool
	got          chan bool

	streams []av.CodecData
}

func NewDemuxer(r io.Reader) *Demuxer {
	d := &Demuxer{
		scanner: bufio.NewScanner(r),
		pkts:    make(chan *av.Packet, 16),
		got:     make(chan bool),
	}

	go func() {
		buf := make([]byte, 0, 64*1024)
		d.scanner.Buffer(buf, 1024*1024)

		// 设置分离函数
		d.scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
			if atEOF && len(data) == 0 {
				return 0, nil, nil
			}

			if i := bytes.Index(data, []byte{0, 0, 0, 1}); i >= 0 {
				return i + 4, data[0:i], nil
			}

			// If we're at EOF, we have a final, non-terminated line. Return it.
			if atEOF {
				return len(data), data, nil
			}
			// Request more data.
			return 0, nil, nil
		})
		// 分离数据
		for d.scanner.Scan() {
			nalu := d.scanner.Bytes()

			if len(nalu) == 0 {
				continue
			}

			p := &av.Packet{
				Idx: 0,
			}

			t := nalu[0] & 0x1f
			if t == h264parser.NALU_PPS {
				d.pps = make([]byte, len(nalu))
				copy(d.pps, nalu)
			} else if t == h264parser.NALU_SPS {
				d.sps = make([]byte, len(nalu))
				copy(d.sps, nalu)
			} else if t != h264parser.NALU_SEI {
				sliceType, err := h264parser.ParseSliceHeaderFromNALU(nalu)
				if err != nil {
					// 合适的错误处理
					fmt.Println(err)
					close(d.pkts)
					return
				}
				if sliceType == h264parser.SLICE_I {
					p.IsKeyFrame = true
				}
			} else {
				continue // NALU_SEI
			}

			// 获取pps和sps
			if d.sps != nil && d.pps != nil && !d.gotCodecData {
				codes, err := h264parser.NewCodecDataFromSPSAndPPS(d.sps, d.sps)
				if err != nil {
					// 合适的错误处理
					fmt.Println(err)
					close(d.got)
					return
				}

				d.streams = append(d.streams, codes)
				d.gotCodecData = true
				d.got <- true
			}
			p.Data = make([]byte, len(nalu))
			copy(p.Data, nalu)
			d.pkts <- p

			log.Printf("read: %0x\n", nalu[:8])

		}
		// 读完了
		close(d.pkts)
	}()
	return d
}

func (d *Demuxer) Streams() (streams []av.CodecData, err error) {
	flag := <-d.got
	if !flag {
		err = fmt.Errorf("Get Streams error")
		return
	}
	streams = d.streams
	return
}

func (d *Demuxer) ReadPacket() (pkt av.Packet, err error) {
	p := <-d.pkts
	if p == nil {
		err = io.EOF
		return
	}
	pkt = *p
	return
}
